import torch
import torch.nn as nn
import torch.utils.model_zoo as model_zoo
from torchvision.models.vgg import vgg16
from torchsummary import torchsummary

class VGG(nn.Module):

    def __init__(self, features, init_weights=False):
        super(VGG, self).__init__()
        self.features = features
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
        if init_weights:
            self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        #x = x.view(x.size(0), -1)
        # x = self.classifier(x)
        #print(x.shape,'final')
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)



def vgg_pool7():

    feature_extract = True
    m_ = vgg16(pretrained=True)
    # torchsummary.summary(model,input_size=(3,224,224))

    layers = torch.nn.Sequential(*list(m_.features))
    model = VGG(layers)
    vgg16_dict = m_.state_dict()
    vgg_dict = {k: v for k, v in vgg16_dict.items() if k in model.state_dict().keys()}
    model.load_state_dict(vgg_dict)
    return model

if __name__ == '__main__':
    model = vgg_pool7()
    torchsummary.summary(model, (3, 224, 224))